import copy
import time
from datetime import datetime
import random
from typing import Any, Dict, Optional
import numpy as np
import gymnasium as gym
from gymnasium import spaces
import cv2
import platform
from scipy.spatial.transform import Rotation as R
import torch

from gym_pih.hardware.RealsenseCamera import Camera as RealsenseCamera
# from gym_pih.hardware.KunWei_ftsensor import ftsensor
from gym_pih.hardware.KWFT import ftsensor
from gym_pih.hardware.pid import PIDForceController
from gym_pih.hardware.LightSource import LightSource
from gym_pih.algorithm.RotationLimiter import RotationLimiter
import rtde_control
import rtde_receive

class TwoStage(gym.Env):
    def __init__(self,
                curriculum=False,
                save_video = False,
                render = False,
                fake_env = False,
                reproduce = False,
                train_algo = 'skrl',
                obs_type = "pixel"):
        self.params_pos_error=0   #孔的定位误差范围
        self.translation_range=0.001   #平移范围1mm
        self.rotation_range=0.04   #旋转范围2.3度
        self.initial_error=0.01   #初始位置与孔的最大距离
        self.ministep_num=6       #每次动作分6步执行

        self.ip="192.168.3.140"
        self.velocity = 0.1
        self.acceleration = 0.05
        self.dt = 1.0/125  # 8ms
        self.lookahead_time = 0.1
        self.gain = 300

        self.obs_type = obs_type
        self.train_algo = train_algo
        self.reproduce = reproduce

        if self.reproduce:
            self.action_space = spaces.Box(low=-1, high=1, shape=(6,), dtype=np.float32)
        else:
            self.action_space = spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32) #3 translation of xyz + 1 rotation along x axis
        if self.obs_type == "stripe":
            self.observation_space = gym.spaces.Dict(
                {"l1":gym.spaces.Box(low=0, high=1, shape=(60,), dtype=np.float32),
                "l2":gym.spaces.Box(low=0,high=1,shape=(60,),dtype=np.float32),
                "l3":gym.spaces.Box(low=0,high=1,shape=(20,),dtype=np.float32),
                "l4":gym.spaces.Box(low=0,high=1,shape=(20,),dtype=np.float32),
                "rec":gym.spaces.Box(low=0,high=1,shape=(225,),dtype=np.float32),
                }
            )
        elif self.obs_type == "pixel":
            if self.train_algo == 'skrl':
                self.observation_space = gym.spaces.Box(low=0, high=1, shape=(64, 64, 3), dtype=np.float32)
            elif self.train_algo == 'sb3':
                self.observation_space = gym.spaces.Box(low=0, high=1, shape=(3, 64, 64), dtype=np.float32)

        self.init_pose=np.zeros(6)
        self.curriculum = curriculum
        self.curriculum_difficulty= 0
        self.success_rate = 0
        self.terminated = False

        
        # assert target_pos.shape==(6,),f"target_pos must be a 6-dim ndarray,but got {target_pos.shape}"
        # self.target_pos=target_pos.copy()

        if self.reproduce:
            self.target_pos=np.array([0.6861132884039939,-0.12243507121158939,0.11622008563921185,0.0001450022568494896,-3.137480067901388,0.0004328617878218785])
        else:
            self.target_pos=np.array([0.6874524902679658, -0.12907513503568732, 0.1543933712663551, -0.00012414404159705842, -3.13757622734271, 3.9352716908090705e-05])
 
        self.pos_error=np.random.uniform(-self.params_pos_error,self.params_pos_error,2)
        self.target_pos[:2]+=self.pos_error#孔的位置有定位误差
        self.pose=np.zeros(6)#用于控制机器人运动

        #运动边界
        self.boundary=np.zeros((2,3))#用于限制机器人的运动范围
        self.boundary[0]=self.target_pos[:3]-0.01
        self.boundary[0][2]=self.target_pos[2]+0.008
        self.boundary[1]=self.target_pos[:3]+0.01
        self.calculate_boundary=self.boundary.copy()
        if self.reproduce:
            self.rotation_limiter = RotationLimiter(self.target_pos[3:6])
        else:
            self.rotation_limiter = RotationLimiter(self.target_pos[3:6],30,0.5,0.5)

        if fake_env:
            return
        # self.Light=LightSource()
        # self.Light.Light_Open(6)
        if self.obs_type == "pixel":
            self.Realsense_cam = RealsenseCamera(640,480)
        elif self.obs_type == "stripe":
            self.Realsense_cam = RealsenseCamera(1280,720)
        self.Realsense_cam.Open()
        self.ft=ftsensor(True)
        self.action_controller=PIDForceController(kp=self.translation_range/self.ministep_num/4,
                                                  ki=self.translation_range/self.ministep_num/100/150,
                                                  kd=self.translation_range/self.ministep_num/4,
                                                  goalForce=-3,downward=False)
        self.rtde_c = rtde_control.RTDEControlInterface(self.ip)
        self.rtde_r = rtde_receive.RTDEReceiveInterface(self.ip)

        self.save_video=save_video
        self.render=render
        if self.save_video:
            self.recording_frames = []

    def _get_obs(self):
        _,wrist_image=self.Realsense_cam.GetImage(isShow=False)
        
        if self.render:
            cv2.namedWindow("wrist",cv2.WINDOW_NORMAL)
            if self.obs_type == "stripe":
                cv2.resizeWindow("wrist",wrist_image.shape[1],wrist_image.shape[0])

                green = (0, 255, 0)
                if self.reproduce:
                    wrist_image[406:466, 654] = green
                    wrist_image[430, 644:664] = green
                    wrist_image[406:466, 693] = green
                    wrist_image[430, 683:703] = green
                    wrist_image[433:448, 660:675] = green
                else:
                    wrist_image[470:530, 645] = green
                    wrist_image[500, 635:655] = green
                    wrist_image[470:530, 688] = green
                    wrist_image[500, 678:698] = green
                    wrist_image[502:517, 660:675] = green
                cv2.imshow("wrist", wrist_image)
            elif self.obs_type == "pixel":
                cv2.resizeWindow("wrist",640,640)
                image = wrist_image[246:310,310:374]
                resized_image = cv2.resize(image, (640, 640), interpolation=cv2.INTER_LINEAR)
                cv2.imshow("wrist", resized_image)

            cv2.waitKey(1)
        if self.save_video:
            self.recording_frames.append(wrist_image)

        if self.obs_type == "stripe":
            gray_image = cv2.cvtColor(wrist_image, cv2.COLOR_BGR2GRAY)
            if self.reproduce:
                l1 = (gray_image[406:466,654]/255.0).astype(np.float32).flatten()
                l4 = (gray_image[430,644:664]/255.0).astype(np.float32).flatten()
                l2 = (gray_image[406:466, 693] / 255.0).astype(np.float32).flatten()
                l3 = (gray_image[430, 683:703] / 255.0).astype(np.float32).flatten()
                rec = gray_image[433:448, 660:675]
            else:
                l1 = (gray_image[470:530, 645] / 255.0).astype(np.float32).flatten()
                l4 = (gray_image[500, 635:655] / 255.0).astype(np.float32).flatten()
                l2 = (gray_image[470:530, 688] / 255.0).astype(np.float32).flatten()
                l3 = (gray_image[500, 678:698] / 255.0).astype(np.float32).flatten()
                rec = gray_image[502:517, 660:675]
            rec = (rec/ 255.0).astype(np.float32)
            rec = rec.flatten()
            observation = {"l1":l1,"l2":l2,"l3":l3,"l4":l4,"rec":rec}
        elif self.obs_type == "pixel":
            wrist_image=wrist_image[246:310,310:374]
            if self.train_algo == 'sb3':
                wrist_image= np.transpose(wrist_image, (2, 0, 1))
            observation = (wrist_image / 255.0).astype(np.float32)
        return observation
    
    def _get_info(self):
        if self.train_algo == "skrl":
            info = {"curriculum":{
                "success_rate": torch.tensor(self.success_rate,dtype = torch.float32),
                "curriculum_difficulty": torch.tensor(self.curriculum_difficulty, dtype = torch.float32)
                 }}
        elif self.train_algo == "sb3":
            info = {"success_rate": self.success_rate,
                    "curriculum_difficulty": self.curriculum_difficulty}
        return info
    
    def _get_reward(self):
        if self.terminated:
            return 1
        return -0.05

    def _generate_random(self,max_value,n):
        if np.random.rand() < 0.5:
            return np.random.uniform(-max_value, -0.001,n)  # 在 [-2, -1] 区间生成
        else:
            return np.random.uniform(0.001, max_value,n)  # 在 [1, 2] 区间生成
    
    def _go_to_reset_pose(self):
        '''
        重置环境，分为是否复现two_stage哪篇论文的实验，若否则做线缆弯曲的重置实验，且由是否terminated决定是否弯曲线缆
        '''
        if self.reproduce:
            transition_pose = self.target_pos.copy()
            transition_pose[2] += 0.01
            self.rtde_c.moveL(transition_pose, 0.2, 0.05)
            
            error_ratio = 0.15+0.35*self.curriculum_difficulty #最低难度误差不为0，避免直接对准孔
            
            init_pose=self.target_pos.copy()
            # init_pose[:2] += np.random.uniform(-self.initial_error*error_ratio,self.initial_error*error_ratio, 2)
            init_pose[:2] += self._generate_random(self.initial_error * error_ratio,2)
            init_pose[2]+=0.01
            rotation = R.from_rotvec(init_pose[3:])
            # Rotations in degrees
            alpha = np.deg2rad(np.random.uniform(-self.curriculum_difficulty, self.curriculum_difficulty))  # Rotation around x-axis
            beta = np.deg2rad(np.random.uniform(-self.curriculum_difficulty, self.curriculum_difficulty))   # Rotation around y-axis
            gamma = np.deg2rad(np.random.uniform(-10*self.curriculum_difficulty, 10*self.curriculum_difficulty))  # Rotation around z-axis

            # Create individual rotations
            rot_x = R.from_rotvec([alpha, 0, 0])
            rot_y = R.from_rotvec([0, beta, 0])
            rot_z = R.from_rotvec([0, 0, gamma])
            final_rotation = rot_z * rot_y * rot_x * rotation
            init_pose[3:] = final_rotation.as_rotvec()
            self.rtde_c.moveL(init_pose, self.velocity, self.acceleration)
        else:
            init_pose=self.target_pos.copy()
            error_ratio = 0.3 + 0.5 * self.curriculum_difficulty
            init_pose[:2] += self._generate_random(self.initial_error * error_ratio,2)
            init_pose[2]+=0.01

            if not self.terminated:
                self.rtde_c.moveL(init_pose, self.velocity, self.acceleration)
                return 0

            if self.curriculum_difficulty==0:
                transition_pose=self.rtde_r.getActualTCPPose()
                transition_pose[2]=init_pose[2]
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                self.rtde_c.moveL(init_pose, 0.2, 0.05)
            elif self.curriculum_difficulty>=0.1 and self.curriculum_difficulty<=0.4:
                transition_pose=self.rtde_r.getActualTCPPose()
                # transition_pose[0]-=0.003
                # self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[2]+=0.002
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1]+=self._randomvariate(0,self.curriculum_difficulty*0.03/0.4)
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                self.rtde_c.moveL(init_pose, 0.2, 0.05)
            else:
                transition_pose=self.rtde_r.getActualTCPPose()
                transition_pose[2]+=0.002
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1]+=0.03
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                down_range = transition_pose[2]-self.calculate_boundary[0][2]+0.01
                transition_pose[2]-=self._randomvariate(0,down_range*(self.curriculum_difficulty-0.4)/0.6)
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                self.rtde_c.moveL(init_pose, 0.2, 0.05)

    def _randomvariate(self,xmin=0, xmax=1):
        """
        Rejection method for generating a single random number from a probability distribution.

        Parameters:
            pdf (function): Probability density function from which you want to generate random numbers.
            xmin (float): Lower bound of the desired range.
            xmax (float): Upper bound of the desired range.

        Returns:
            float: A random number following the input probability distribution.
        """

        def pdf(x):
            return abs(x)*2

        # Calculate the minimal and maximum values of the PDF in the desired interval.
        # The rejection method needs these values in order to work properly.
        # x = np.linspace(xmin, xmax, 1000)
        # y = pdf(x)
        # pmin = 0.
        # pmax = y.max()

        #这里的pdf是一个线性函数，所以pmin=0，pmax=pdf(xmax)
        pmin = 0
        pmax = pdf(xmax)

        # Keep generating numbers until we accept one
        while True:
            x_prime = np.random.uniform(xmin, xmax)  # x'
            y_prime = np.random.uniform(pmin, pmax)  # y'

            if y_prime < pdf(x_prime):
                return x_prime

    def _miniServoL(self, action, step=4, is_insertion = True):
        """
        输入一次动作，分多步执行

        :param action: 动作
        :param step: 步数，默认为4
        :parm is_insertion: 如果要执行_insertion_trial()则不受运动边界限制
        """
        miniaction = action / step
        for i in range(step):
            force=self.ft.GetForce()
            
            delta=self.action_controller._calPIDForce(force)
            self.pose[:3] += miniaction[:3]+delta[:3]
            # if abs(force[2])>5:
            #     print("force=",force)
            current_rotation = R.from_rotvec(self.pose[3:])
            delta_rotation = R.from_rotvec(miniaction[3:])
            new_rotation = (current_rotation * delta_rotation).as_rotvec()
            current_vec = self.rotation_limiter.apply_limits(new_rotation)
            self.pose[3:] = current_vec
            # self.pose[3:] = new_rotation
            if not is_insertion:
                self.pose[:3]=np.clip(self.pose[:3],self.boundary[0],self.boundary[1])
            t_start=self.rtde_c.initPeriod()
            self.rtde_c.servoL(self.pose, self.velocity, self.acceleration, self.dt, self.lookahead_time, self.gain)
            self.rtde_c.waitPeriod(t_start)

    def _insertion_trail(self):
        '''
        原文中的ap，执行每个动作后尝试向下插入
        '''
        force=self.ft.GetForce()
        while abs(force[2])<5 and self.pose[2]>self.target_pos[2]-0.003:
            delta = self.pose[2]-self.target_pos[2]+0.05
            if delta>0.001:
                action = np.array([0, 0, -0.001, 0, 0, 0])
                self._miniServoL(action, 10)
                force=self.ft.GetForce()
                terminated = self.rtde_r.getDigitalInState(0)
                if terminated:
                    return True
            else:
                action = np.array([0]*6)
                action[2] -= delta
                self._miniServoL(action, 10)
                terminated = self.rtde_r.getDigitalInState(0)
                if terminated:
                    return True
        terminated = self.rtde_r.getDigitalInState(0)
        return terminated

    def _update_boundary(self):

        if self.reproduce:
            rotation_ratio = 2+8*self.curriculum_difficulty
            ratio = 0.2 + 0.3 * self.curriculum_difficulty
            self.rotation_limiter.update_limit(max_angle_z_deg=rotation_ratio)
        else:
            rotation_ratio = 2+18*self.curriculum_difficulty
            ratio = 0.4 + 0.4 * self.curriculum_difficulty
            self.rotation_limiter.update_limit(max_angle_x_deg=rotation_ratio)

        self.boundary[0] = self.target_pos[:3] - 0.01 * ratio
        self.boundary[0][2] = self.target_pos[2] + 0.008
        self.boundary[1] = self.target_pos[:3] + 0.01 * ratio
        self.boundary[1][2] = self.target_pos[2] + 0.01
              
    def reset(self, seed: int = None, options: Optional[Dict[str, Any]] = None):
        # if not self.curriculum:
        #     self.curriculum_difficulty=random.uniform(0.0,1.0)

        self._update_boundary()
            
        self.rtde_c.servoStop()        #ur进入servo模式后，必须先停止才能再次用move
        self._go_to_reset_pose()

        self.pose=self.rtde_r.getActualTCPPose()
        self.init_pose=self.pose.copy()
        observation=self._get_obs()
        info=self._get_info()
        self.terminated = False
        
        return observation,info
    
    def step(self, action):
        if self.reproduce:
            execute_action = action.copy()
            rotation_ratio = np.array([0.1 * self.rotation_range, 0.1 * self.rotation_range, self.rotation_range])
            execute_action[3:] = execute_action[3:] * rotation_ratio
        else:
            execute_action = np.zeros(6)
            execute_action[3] = action[3]*self.rotation_range/2
        execute_action[:3] = action[:3] * self.translation_range

        self._miniServoL(execute_action,self.ministep_num,False)
        observation=self._get_obs()
        init_pose = self.pose.copy()
        self.terminated=self._insertion_trail()
        if not self.terminated:
            self.rtde_c.servoStop()
            self.pose = init_pose
            self.rtde_c.moveL(self.pose, self.velocity, self.acceleration)
        reward=self._get_reward()
        info=self._get_info()

        return observation, reward, self.terminated, False, info
    
    def close(self):
        self.rtde_c.servoStop()
        self.rtde_c.stopScript()
        self.rtde_c.disconnect()
        self.rtde_r.disconnect()
        self.ft.DisConnect()
        # self.Light.Light_Close()
        # self.Light.port_close()
